from flask import g
from functools import wraps

from models import db


def set_db_to_read(func):
    """
    设置使用读数据库
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        db.session().set_to_read()
        return func(*args, **kwargs)
    return wrapper


def set_db_to_write(func):
    """
    设置使用写数据库
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        db.session().set_to_write()
        return func(*args, **kwargs)
    return wrapper


def login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        # 如果用户已登录, 并且发送的是访问token, 允许访问视图
        if g.user_id and g.is_refresh == False:
            return f(*args, **kwargs)
        else:  # 失败, 返回401状态
            return {'message': 'Invalid token'}, 401

    return wrapper